#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time   : 2023-06-27 23:11
# @Author : hard work
# @File    : snmp_trap_util.py
# @Description :
"""ok"""

from pysnmp.carrier.asynsock.dispatch import AsynsockDispatcher
from pysnmp.carrier.asynsock.dgram import udp
from pyasn1.codec.ber import decoder
from pysnmp.proto import api
from pysnmp.proto.rfc1905 import VarBind
import re

from test_app_main.case_cfg.snmp_oid_cfg import nari_check_trap_entry
from utils.timeControl import now_time
from utils.get_local_ip import get_host_ip

def pick(varbind):
    pattern = r"name=(.*)[\s\S]*value=(.*)"
    regx = re.compile(pattern)
    matchs = regx.findall(varbind)
    if matchs:
        return matchs[0][0], matchs[0][1]
    else:
        return False


def cbFun(transportDispatcher, transportDomain, transportAddress, wholeMsg):
    # nari_check_dict = {}
    while wholeMsg:
        msgVer = int(api.decodeMessageVersion(wholeMsg))
        if msgVer in api.protoModules:
            pMod = api.protoModules[msgVer]
        else:
            print('Unsupported SNMP version %s' % msgVer)
            return
        reqMsg, wholeMsg = decoder.decode(wholeMsg, asn1Spec=pMod.Message(), )
        # print("reqmsg:", reqMsg)
        # print("wholemsg:",wholeMsg)
        print('Notification message from %s:%s: ' % (transportDomain, transportAddress))
        reqPDU = pMod.apiMessage.getPDU(reqMsg)
        # print("pdu:",reqPDU)
        varBinds = pMod.apiPDU.getVarBindList(reqPDU)
        # print("varbinds:",varBinds)
        for row in varBinds:
            row: VarBind
            row = row.prettyPrint()
            k, v = pick(row)
            for mib_name, oid in nari_check_trap_entry().items():
                # print(k1, v1)
                if oid in k:
                    print(mib_name, ": " + oid, "=> " + v)
                    # nari_check_dict.update({oid: v})
                    print(now_time(), 100 * "=")
                    # print("%s:%s" % (k, v))
                else:
                    # print(mib_name, ": " + oid, "!=> " + v)
                    pass
            # print("%s:%s" % (k, v))
    return wholeMsg


def snmp_trap_main(client_ip: str, port: int):
    transportDispatcher = AsynsockDispatcher()
    transportDispatcher.registerRecvCbFun(cbFun)
    # UDP/IPv4
    transportDispatcher.registerTransport(
        udp.domainName, udp.UdpSocketTransport().openServerMode((client_ip, port))
    )
    # UDP/IPv6
    # transportDispatcher.registerTransport(
    #     udp6.domainName, udp6.Udp6SocketTransport().openServerMode(('::1', 162))
    # )
    transportDispatcher.jobStarted(1)
    try:
        transportDispatcher.runDispatcher()

    except KeyboardInterrupt:
        raise
    except:
        transportDispatcher.closeDispatcher()
        raise


if __name__ == '__main__':

    print(get_host_ip())
    port = 162
    snmp_trap_main(get_host_ip(), port)
